package org.laokouyun.demo;

import io.grpc.StatusRuntimeException;
import io.grpc.stub.AbstractBlockingStub;
import lombok.RequiredArgsConstructor;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.loadbalancer.LoadBalancerClient;
import org.springframework.grpc.client.GrpcClientFactory;
import org.springframework.stereotype.Component;
import org.springframework.util.ReflectionUtils;

import java.lang.reflect.Method;
import java.time.Duration;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Component
@RequiredArgsConstructor
public final class GpcClientUtils {

    private final Map<String, Object> stubMap = new ConcurrentHashMap<>();

    private final Map<String, Method> methodMap = new ConcurrentHashMap<>();

    private final LoadBalancerClient loadBalancerClient;

    private final GrpcClientFactory grpcClientFactory;

    private final int maxRetry = 3;

    private final int waitMillis = 100;

    public <T extends AbstractBlockingStub<T>, P> Object invoke(String serviceId, Class<T> clazz, String methodName, Class<P> paramClazz, Object...args) throws Exception {
        Method method = getMethod(clazz, methodName, paramClazz);
        for (int attempt  = 1; attempt  <= maxRetry; attempt++) {
            try {
                return ReflectionUtils.invokeMethod(method, getStub(serviceId, clazz), args);
            } catch (Exception ex) {
                if (!(ex instanceof StatusRuntimeException)) {
                    throw ex;
                } else {
                    Thread.sleep(Duration.ofMillis(waitMillis));
                }
            }
        }
        return ReflectionUtils.invokeMethod(method, FallbackFactory.getFallback(clazz), args);
    }

    @SuppressWarnings("unchecked")
    private <T extends AbstractBlockingStub<T>> T getStub(String serviceId, Class<T> clazz) {
        ServiceInstance serviceInstance = loadBalancerClient.choose(serviceId);
        if (serviceInstance == null) {
            throw new IllegalStateException(serviceId + " is not available");
        }
        String target = String.format("static://%s:%s", serviceInstance.getHost(), serviceInstance.getMetadata().getOrDefault("grpc_port", "9090"));
        String key = String.format("%s:%s:%s", serviceId, clazz.getName(), target);
        return (T) stubMap.computeIfAbsent(key, _ -> grpcClientFactory.getClient(target, clazz, null));
    }

    private  <T extends AbstractBlockingStub<T>, P> Method getMethod(Class<T> clazz, String methodName, Class<P> paramClazz) {
        String key = String.format("%s:%s:%s", clazz.getName(), methodName, paramClazz.getName());
        return methodMap.computeIfAbsent(key, _ -> {
            Method method = ReflectionUtils.findMethod(clazz, methodName, paramClazz);
            assert method != null;
            ReflectionUtils.makeAccessible(method);
            return method;
        });
    }

}
